import os
import qianfan

# 通过环境变量初始化认证信息
# 方式一：【推荐】使用安全认证AK/SK
# 替换下列示例中参数，安全认证Access Key替换your_iam_ak，Secret Key替换your_iam_sk，如何获取请查看https://cloud.baidu.com/doc/Reference/s/9jwvz2egb
os.environ["QIANFAN_ACCESS_KEY"] = os.getenv('QIANFAN_ACCESS_KEY', '')
os.environ["QIANFAN_SECRET_KEY"] = os.getenv('QIANFAN_SECRET_KEY', '')

# 方式二：【不推荐】使用应用AK/SK鉴权
# 替换下列示例中参数，将应用API_Key、应用Secret key值替换为真实值
# os.environ["QIANFAN_AK"] = os.getenv('QIANFAN_AK', '')
# os.environ["QIANFAN_SK"] = os.getenv('QIANFAN_SK', '')

# 初始化聊天完成对象
chat_comp = qianfan.ChatCompletion()


def send_message(chat_comp, model, messages):
    """
    发送消息并获取响应。

    Args:
        chat_comp (qianfan.ChatCompletion): 聊天完成对象。
        model (str): 使用的模型名称。
        messages (list): 对话消息列表。

    Returns:
        str: 助手的完整回复。
    """
    # 发送请求并获取响应
    resp = chat_comp.do(
        model=model,
        messages=messages
    )
    # 返回助手的回复
    return resp["body"]['result']


def main():
    """
    主函数。
    """
    # 设置模型名称
    model = "ERNIE-Lite-8K"

    # 初始对话历史
    messages = [
        {
            "role": "system",
            "content": "你是一个有趣的SQL助手，可以回答任何问题。",
        }, {
            "role": "user",
            "content": "你好",
        }
    ]

    while True:
        # 发送消息并获取响应
        response_content = send_message(chat_comp, model, messages)

        # 将响应添加到对话历史
        messages.append({"role": "assistant", "content": response_content})

        # 打印助手的回复
        print(f"助手: {response_content}")

        # 获取用户输入
        user_input = input("\n你: ")
        # 如果用户输入 "exit" 或 "quit"，退出循环
        if user_input.lower() in ["exit", "quit"]:
            break

        # 将用户输入添加到对话历史
        messages.append({"role": "user", "content": user_input})


if __name__ == "__main__":
    main()
